#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2021/6/29 22:19
# @Author  : LiShan
# @Email   : lishan_1997@126.com
# @File    : flowrate_count.py
# @Note    : this is note
import os
from shutil import move
import cv2.cv2 as cv
import time

save_path = "./flow"
width = 480
height = 360

path_list = [
    save_path,
    save_path + "/fram",
    save_path + "/bg",
    save_path + "/fgmask",
    save_path + "/fg",
]

for path in path_list:
    if os.path.exists(path):
        pass
    else:
        os.mkdir(path)


class TrafficFlowDetected:
    def __init__(self):
        self.video = "./video/test.mp4"
        self.width = 640
        self.height = 480

    # 目标预检测
    def pre_object_detect(self, contours):
        pre_object = []
        for cnt in contours:
            # 依次取出每一条轮廓，计算点集或灰度图像的非零像素的右上边界矩形
            x, y, w, h = cv.boundingRect(cnt)
            rect = x, y, w, h
            area = self.width * self.height
            # 过滤掉小物体、大物体
            c1 = cv.contourArea(cnt) < area * 0.01 or cv.contourArea(cnt) > area * 0.25
            c2 = w < self.width * 0.05 or w > self.width * 0.25
            c3 = h < self.height * 0.05 or h > self.height * 0.25
            if c1 or c2 or c3:
                continue
            else:
                pre_object.append(rect)
        return pre_object

    # 判断目标框位置，删除内部框
    def is_inside(self, o, i):
        # o为第一个框，i为第二个框
        ox, oy, ow, oh = o
        ix, iy, iw, ih = i
        # 如果第一个矩形框被完全包含在第二个矩形框中，可确定第一个矩形框应该被丢弃
        return ox > ix and oy > iy and ox + ow < ix + iw and oy + oh < iy + ih

    # 遍历检测结果来丢弃不含有检测目标的区域
    def dropout(self, pre_object):
        object = []
        num = 0
        # ri、qi为框的编号，r、q为对应框的x，y，w，h坐标参数
        for ri, r in enumerate(pre_object):
            for qi, q in enumerate(pre_object):
                if ri != qi and self.is_inside(r, q):
                    break
                else:
                    num += 1
            if num == len(pre_object):
                object.append(r)
            num = 0
        return object

    # 目标检测
    def target_detected(self, video=None):
        if video is None:
            video = self.video
        mask = cv.imread("./lane" + "/mask.jpg")
        new_mask = cv.imread("./lane" + "/new_mask.jpg")
        # new_mask = cv.bitwise_not(new_mask)
        cap = cv.VideoCapture(video)
        fps = cap.get(cv.CAP_PROP_FPS)

        # 背景差分法（混合高斯建模）创建一个对象
        subtractor = cv.createBackgroundSubtractorMOG2(detectShadows=False)
        frame_num = 0
        try:
            while cap.isOpened():
                ret, frame = cap.read()
                if not ret:
                    pass
                else:
                    frame_num += 1
                    # 调整图片大小：原图、（宽，高）
                    frame = cv.resize(frame, (self.width, self.height))
                    cv.imwrite(save_path + "/fram/" + str(frame_num) + "_frame" + ".jpg", frame)
                    # 原图
                    origin = frame
                    # 感兴趣区域提取
                    roi = cv.bitwise_and(new_mask, frame)
                    cv.imwrite(save_path + "/roi/" + str(frame_num) + "_roi" + ".jpg", roi)
                    # 转为灰度图
                    gray = cv.cvtColor(roi, cv.COLOR_BGR2GRAY)
                    cv.imwrite(save_path + "/gray/" + str(frame_num) + "_gray" + ".jpg", gray)
                    # 根据背景差分法提取前景腌膜
                    fgmask = subtractor.apply(gray)
                    # new_fgmask = cv.bitwise_not(fgmask)
                    cv.imwrite(save_path + "/fgmask/" + str(frame_num) + "_fgmask" + ".jpg", fgmask)
                    # 二值化
                    # ret, binary = cv.threshold(fgmask, 50, 255, cv.THRESH_BINARY)
                    # binary = cv.adaptiveThreshold(fgmask, 255, cv.ADAPTIVE_THRESH_MEAN_C, cv.THRESH_BINARY, 11, 5)
                    binary = cv.adaptiveThreshold(fgmask, 255, cv.ADAPTIVE_THRESH_GAUSSIAN_C, cv.THRESH_BINARY, 11, 5)
                    binary = cv.bitwise_not(binary)
                    cv.imwrite(save_path + "/binary/" + str(frame_num) + "_binary" + ".jpg", binary)
                    # 形态学操作
                    element = cv.getStructuringElement(cv.MORPH_RECT, (3, 3))
                    element2 = cv.getStructuringElement(cv.MORPH_RECT, (3, 3))
                    # 开运算
                    open = cv.morphologyEx(binary, cv.MORPH_OPEN, element)
                    cv.imwrite(save_path + "/open/" + str(frame_num) + "_open" + ".jpg", open)
                    # 膨胀运算
                    inflation = cv.dilate(open, element2)
                    cv.imwrite(save_path + "/inflation/" + str(frame_num) + "_inflation" + ".jpg", inflation)

                    # 二值图像提取轮廓，轮廓跟踪算法（Suzuki，1985）
                    contours, hierarchy = cv.findContours(inflation, cv.RETR_TREE, cv.CHAIN_APPROX_NONE)
                    contour_origin = origin.copy()
                    cnt = contours[0:len(contours) - 2]
                    cv.drawContours(contour_origin, cnt, -1, (0, 0, 255), 2)
                    cv.imwrite(save_path + "/contour/" + str(frame_num) + "_contour" + ".jpg", contour_origin)
                    pre_object = self.pre_object_detect(contours)
                    object = self.dropout(pre_object)
                    for i in range(len(object)):
                        x, y, w, h = object[i]
                        cv.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
                    cv.imwrite(save_path + "/object/" + str(frame_num) + "_object" + ".jpg", frame)
                    k = cv.waitKey(1) & 0xff
                    if k == 27:
                        break
            cap.release()
            cv.destroyAllWindows()
        except:
            pass


if __name__ == '__main__':
    TrafficFlowDetected().target_detected("./video/background_model.mp4")
